"""
直接流域划分工具
提供一个简单的函数，直接输入坐标就能输出shp文件，无需创建CSV文件

使用示例:
    from direct_delineate import delineate_direct
    
    # 直接输入坐标获取流域边界并保存为shp文件
    delineate_direct(
        lat=64.71072,
        lng=-21.60337,
        outlet_name="Nordhura River at Stekkur",
        outlet_id="6401070",
        output_path="watershed.shp",
        high_res=True,
        area_estimate=507
    )
"""

import os
import sys
import tempfile
import shutil

import geopandas as gpd
import pandas as pd
from loguru import logger

# 使用相对导入
from ..core import delineate
from ..utils import config


def delineate_direct(lat, lng, outlet_name="Custom_Outlet", outlet_id="1",
                     output_path="watershed.geojson", high_res=True, area_estimate=None,
                     output_format="geojson", copy_to_final=True):
    """
    直接使用坐标进行流域划分，无需创建CSV文件
    
    参数:
        lat: 纬度
        lng: 经度
        outlet_name: 出口点名称
        outlet_id: 出口点ID
        output_path: 输出文件路径
        high_res: 是否使用高分辨率模式
        area_estimate: 预估面积（可选）
        output_format: 输出格式 ('shp', 'gpkg', 'geojson')
        copy_to_final: 是否复制到最终输出目录（如果配置了FINAL_OUTPUT_DIR）
    
    返回:
        成功时返回输出文件路径，失败时返回None
    """
    import config

    # 创建临时目录
    temp_dir = tempfile.mkdtemp()
    temp_csv = os.path.join(temp_dir, 'temp_outlets.csv')

    # 保存原始配置
    original_outlets_csv = config.OUTLETS_CSV
    original_high_res = config.HIGH_RES
    original_output_dir = config.OUTPUT_DIR
    original_output_ext = config.OUTPUT_EXT
    original_verbose = config.VERBOSE

    try:
        # 创建出口点数据
        outlet_data = {
            'id': [outlet_id],
            'lat': [lat],
            'lng': [lng],
            'name': [outlet_name]
        }

        if area_estimate:
            outlet_data['area'] = [area_estimate]

        # 创建临时CSV文件
        df = pd.DataFrame(outlet_data)
        df.to_csv(temp_csv, index=False)

        logger.info(f"创建临时CSV文件: {temp_csv}")
        logger.info(f"出口点坐标: ({lat}, {lng})")
        logger.info(f"出口点名称: {outlet_name}")
        if area_estimate:
            logger.info(f"预估面积: {area_estimate} km²")

        # 设置环境变量（这是config.py读取配置的方式）
        os.environ['OUTLETS_CSV'] = temp_csv
        os.environ['HIGH_RES'] = str(high_res)
        os.environ['OUTPUT_DIR'] = config.OUTPUT_DIR  # 使用配置中的输出目录，而不是临时目录
        os.environ['VERBOSE'] = 'True'
        os.environ['PLOTS'] = str(config.PLOTS).lower()  # 添加PLOTS环境变量

        # 根据输出格式设置文件扩展名
        if output_format.lower() == 'shp':
            os.environ['OUTPUT_EXT'] = "shp"
        elif output_format.lower() == 'gpkg':
            os.environ['OUTPUT_EXT'] = "gpkg"
        elif output_format.lower() == 'geojson':
            os.environ['OUTPUT_EXT'] = "geojson"
        else:
            raise ValueError(f"不支持的输出格式: {output_format}")

        # 重新导入config以获取更新的配置
        import importlib
        importlib.reload(config)

        # 由于delineate.py使用了from config import *，我们需要重新导入delineate模块
        # 以确保它使用更新后的配置
        import delineate
        importlib.reload(delineate)

        # 执行流域划分
        logger.info("开始执行流域划分...")
        logger.info(f"使用CSV文件: {config.OUTLETS_CSV}")
        logger.info(f"输出目录: {config.OUTPUT_DIR}")
        logger.info(f"输出格式: {config.OUTPUT_EXT}")
        delineate.delineate()

        # 查找输出文件 - delineate函数输出到config.OUTPUT_DIR
        output_files = []
        output_dir = config.OUTPUT_DIR  # 使用配置中的输出目录
        logger.info(f"在目录中查找输出文件: {output_dir}")

        # 根据输出格式确定要查找的文件扩展名
        if output_format.lower() == 'geojson':
            target_extensions = ['.geojson']
        elif output_format.lower() == 'shp':
            target_extensions = ['.shp']
        elif output_format.lower() == 'gpkg':
            target_extensions = ['.gpkg']
        else:
            target_extensions = ['.gpkg', '.shp', '.geojson']  # 兼容所有格式

        # 确保输出目录存在
        if not os.path.exists(output_dir):
            logger.warning(f"输出目录不存在，尝试在主输出目录中查找: {config.BASE_OUTPUT_DIR}")
            # 如果配置的输出目录不存在，尝试在主输出目录中查找
            if os.path.exists(config.BASE_OUTPUT_DIR):
                output_dir = config.BASE_OUTPUT_DIR
            else:
                logger.error(f"主输出目录也不存在: {config.BASE_OUTPUT_DIR}")
                return None

        if os.path.exists(output_dir):
            for file in os.listdir(output_dir):
                logger.info(f"检查文件: {file}")
                # 检查文件是否以outlet_id开头并且有正确的扩展名
                if file.startswith(outlet_id):
                    for ext in target_extensions:
                        if file.endswith(ext):
                            output_files.append(os.path.join(output_dir, file))
                            logger.info(f"找到匹配文件: {file}")
                            break

        if not output_files:
            logger.error(
                f"未找到输出的流域边界文件，查找条件: 以 {outlet_id} 开头的 {' 或 '.join(target_extensions)} 文件")
            logger.info(f"输出目录内容: {os.listdir(output_dir) if os.path.exists(output_dir) else '目录不存在'}")
            return None

        # 读取结果并保存到指定路径
        if len(output_files) == 1:
            # 读取文件并转换为指定格式
            gdf = gpd.read_file(output_files[0])

            # 保存为指定格式
            if output_format.lower() == 'shp':
                gdf.to_file(output_path)
            elif output_format.lower() == 'gpkg':
                gdf.to_file(output_path, driver='GPKG')
            elif output_format.lower() == 'geojson':
                gdf.to_file(output_path, driver='GeoJSON')

            logger.success(f"流域边界已保存到: {output_path}")
            
            # 复制到最终输出目录（如果配置了且启用了复制）
            if copy_to_final and config.FINAL_OUTPUT_DIR and config.FINAL_OUTPUT_DIR.strip():
                final_path = _copy_to_final_output(output_path, outlet_id, config.FINAL_OUTPUT_DIR.strip())
                if final_path:
                    logger.info(f"文件已复制到最终输出目录: {final_path}")
            
            return output_path
        else:
            # 合并多个文件
            all_watersheds = []
            for file in output_files:
                gdf = gpd.read_file(file)
                all_watersheds.append(gdf)

            combined_gdf = gpd.GeoDataFrame(pd.concat(all_watersheds, ignore_index=True))

            # 保存合并后的结果
            if output_format.lower() == 'shp':
                combined_gdf.to_file(output_path)
            elif output_format.lower() == 'gpkg':
                combined_gdf.to_file(output_path, driver='GPKG')
            elif output_format.lower() == 'geojson':
                combined_gdf.to_file(output_path, driver='GeoJSON')

            logger.success(f"合并后的流域边界已保存到: {output_path}")
            
            # 复制到最终输出目录（如果配置了且启用了复制）
            if copy_to_final and config.FINAL_OUTPUT_DIR and config.FINAL_OUTPUT_DIR.strip():
                final_path = _copy_to_final_output(output_path, outlet_id, config.FINAL_OUTPUT_DIR.strip())
                if final_path:
                    logger.info(f"文件已复制到最终输出目录: {final_path}")
            
            return output_path

    except Exception as e:
        logger.error(f"流域划分失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

    finally:
        # 恢复原始配置
        config.OUTLETS_CSV = original_outlets_csv
        config.HIGH_RES = original_high_res
        config.OUTPUT_DIR = original_output_dir
        config.OUTPUT_EXT = original_output_ext
        config.VERBOSE = original_verbose

        # 清理临时文件
        import shutil
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)


def _copy_to_final_output(source_path: str, outlet_id: str, final_dir: str) -> str:
    """
    复制文件到最终输出目录
    
    参数:
        source_path: 源文件路径
        outlet_id: 出口点ID
        final_dir: 最终输出目录
    
    返回:
        成功时返回最终文件路径，失败时返回None
    """
    try:
        # 确保最终输出目录存在
        os.makedirs(final_dir, exist_ok=True)
        
        # 生成目标文件路径
        filename = os.path.basename(source_path)
        dest_path = os.path.join(final_dir, filename)
        
        # 复制文件
        shutil.copy2(source_path, dest_path)
        logger.info(f"文件已复制: {source_path} -> {dest_path}")
        
        return dest_path
        
    except Exception as e:
        logger.error(f"复制文件到最终输出目录失败: {str(e)}")
        return None


def main():
    """
    主函数，演示如何使用直接流域划分功能
    """
    logger.info("=== 直接流域划分工具演示 ===")

    # # 冰岛 Nordhura River at Stekkur
    # result = delineate_direct(
    #     lat=64.71072,
    #     lng=-21.60337,
    #     outlet_name="Nordhura River at Stekkur",
    #     outlet_id="6401070",
    #     output_path="nordhura_watershed.geojson",
    #     high_res=True,
    #     area_estimate=507,
    #     output_format="geojson"
    # )

    # 清河水库
    result = delineate_direct(
        lat=42.53819,
        lng=124.1693,
        outlet_name="清河水库",
        outlet_id="qinghe_reservoir",
        output_path="qinghe_watershed.geojson",
        high_res=True,
        area_estimate=2379.852,
        output_format="geojson"
    )

    # # 大伙房水库
    # result = delineate_direct(
    #     lat=41.883333,
    #     lng=124.083333,
    #     outlet_name="大伙房水库",
    #     outlet_id="21100150",
    #     output_path="21100150.geojson",
    #     high_res=True,
    #     area_estimate=5471.007,
    #     output_format="geojson"
    # )

    if result:
        logger.success(f"成功！流域边界已保存到: {result}")
    else:
        logger.error("流域划分失败")


if __name__ == "__main__":
    main()
